import os
import sys
import json
import shutil
import requests
import tempfile
import gradio as gr
import pandas as pd

from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm


now_dir = os.getcwd()
sys.path.append(now_dir)

from core import run_download_script
from rvc.lib.utils import format_title

from assets.i18n.i18n import I18nAuto

i18n = I18nAuto()

gradio_temp_dir = os.path.join(tempfile.gettempdir(), "gradio")

if os.path.exists(gradio_temp_dir):
    shutil.rmtree(gradio_temp_dir)


def save_drop_model(dropbox):
    if "pth" not in dropbox and "index" not in dropbox:
        raise gr.Error(
            message="The file you dropped is not a valid model file. Please try again."
        )

    file_name = format_title(os.path.basename(dropbox))
    model_name = file_name

    if ".pth" in model_name:
        model_name = model_name.split(".pth")[0]
    elif ".index" in model_name:
        replacements = ["nprobe_1_", "_v1", "_v2", "added_"]
        for rep in replacements:
            model_name = model_name.replace(rep, "")
        model_name = model_name.split(".index")[0]

    model_path = os.path.join(now_dir, "logs", model_name)
    if not os.path.exists(model_path):
        os.makedirs(model_path)
    if os.path.exists(os.path.join(model_path, file_name)):
        os.remove(os.path.join(model_path, file_name))
    shutil.move(dropbox, os.path.join(model_path, file_name))
    print(f"{file_name} saved in {model_path}")
    gr.Info(f"{file_name} saved in {model_path}")

    return None


json_url = "https://huggingface.co/IAHispano/Applio/raw/main/pretrains.json"


def fetch_pretrained_data():
    pretraineds_custom_path = os.path.join("rvc", "models", "pretraineds", "custom")
    os.makedirs(pretraineds_custom_path, exist_ok=True)
    try:
        with open(
            os.path.join(pretraineds_custom_path, json_url.split("/")[-1]), "r"
        ) as f:
            data = json.load(f)
    except:
        try:
            response = requests.get(json_url)
            response.raise_for_status()
            data = response.json()
            with open(
                os.path.join(pretraineds_custom_path, json_url.split("/")[-1]),
                "w",
                encoding="utf-8",
            ) as f:
                json.dump(
                    data,
                    f,
                    indent=2,
                    separators=(",", ": "),
                    ensure_ascii=False,
                )
        except:
            data = {
                "Titan": {
                    "32k": {"D": "null", "G": "null"},
                },
            }
    return data


def get_pretrained_list():
    data = fetch_pretrained_data()
    return list(data.keys())


def get_pretrained_sample_rates(model):
    data = fetch_pretrained_data()
    return list(data[model].keys())


def get_file_size(url):
    response = requests.head(url)
    return int(response.headers.get("content-length", 0))


def download_file(url, destination_path, progress_bar):
    os.makedirs(os.path.dirname(destination_path), exist_ok=True)
    response = requests.get(url, stream=True)
    block_size = 1024
    with open(destination_path, "wb") as file:
        for data in response.iter_content(block_size):
            file.write(data)
            progress_bar.update(len(data))


def download_pretrained_model(model, sample_rate, url_g="", url_d=""):
    save_path = os.path.join("rvc", "models", "pretraineds", "custom")
    os.makedirs(save_path, exist_ok=True)
    tasks = []

    if url_g or url_d:
        tasks = [
            (u, os.path.join(save_path, os.path.basename(u)))
            for u in [url_g, url_d]
            if u
        ]
        if not tasks:
            return gr.Warning(i18n("Please provide at least one URL."))
    else:
        data = fetch_pretrained_data()
        paths = data[model][sample_rate]
        tasks = [
            (
                f"https://huggingface.co/{p}",
                os.path.join(save_path, os.path.basename(p)),
            )
            for p in [paths["D"], paths["G"]]
        ]

    gr.Info(i18n("Downloading pretrained model..."))

    with tqdm(
        total=sum(get_file_size(u) for u, _ in tasks),
        unit="iB",
        unit_scale=True,
        desc="Downloading files",
    ) as pbar:
        with ThreadPoolExecutor(max_workers=2) as executor:
            futures = [
                executor.submit(download_file, url, dst, pbar) for url, dst in tasks
            ]
            for f in futures:
                f.result()

    gr.Info(i18n("Pretrained model downloaded successfully!"))
    print("Pretrained model downloaded successfully!")


def update_sample_rate_dropdown(model):
    return {
        "choices": get_pretrained_sample_rates(model),
        "value": get_pretrained_sample_rates(model)[0],
        "__type__": "update",
    }


def download_handler(is_custom, model, sample_rate, url_g, url_d):
    if is_custom:
        download_pretrained_model(
            None,
            None,
            url_g.replace("?download=true", ""),
            url_d.replace("?download=true", ""),
        )
    else:
        download_pretrained_model(model, sample_rate, "", "")


def download_tab():
    with gr.Column():
        gr.Markdown(value=i18n("## Download Model"))
        model_link = gr.Textbox(
            label=i18n("Model Link"),
            placeholder=i18n("Introduce the model link"),
            interactive=True,
        )
        model_download_output_info = gr.Textbox(
            label=i18n("Output Information"),
            info=i18n("The output information will be displayed here."),
            value="",
            max_lines=8,
            interactive=False,
        )
        model_download_button = gr.Button(i18n("Download Model"))
        model_download_button.click(
            fn=run_download_script,
            inputs=[model_link],
            outputs=[model_download_output_info],
        )
        gr.Markdown(value=i18n("## Drop files"))
        dropbox = gr.File(
            label=i18n(
                "Drag your .pth file and .index file into this space. Drag one and then the other."
            ),
            type="filepath",
        )

        dropbox.upload(
            fn=save_drop_model,
            inputs=[dropbox],
            outputs=[dropbox],
        )
        gr.Markdown(value=i18n("## Download Pretrained Models"))

        with gr.Group():
            with gr.Group(visible=True) as default:
                pretrained_model = gr.Dropdown(
                    label=i18n("Pretrained"),
                    info=i18n("Select the pretrained model you want to download."),
                    choices=get_pretrained_list(),
                    value="Titan",
                    interactive=True,
                )
                pretrained_sample_rate = gr.Dropdown(
                    label=i18n("Sampling Rate"),
                    info=i18n("And select the sampling rate."),
                    choices=get_pretrained_sample_rates(pretrained_model.value),
                    value="40k",
                    interactive=True,
                    allow_custom_value=True,
                )

            with gr.Group(visible=False) as custom:
                url_g = gr.Textbox(label=i18n("Pretrained G"), interactive=True)
                url_d = gr.Textbox(label=i18n("Pretrained D"), interactive=True)

            use_custom = gr.Checkbox(
                label=i18n("Custom Pretrained"),
                value=False,
                interactive=True,
            )

        download_pretrained = gr.Button(i18n("Download"))

        pretrained_model.change(
            update_sample_rate_dropdown,
            inputs=[pretrained_model],
            outputs=[pretrained_sample_rate],
        )

        use_custom.change(
            fn=lambda x: (
                {"visible": not x, "__type__": "update"},
                {"visible": x, "__type__": "update"},
            ),
            inputs=[use_custom],
            outputs=[default, custom],
        )

        download_pretrained.click(
            fn=download_handler,
            inputs=[use_custom, pretrained_model, pretrained_sample_rate, url_g, url_d],
            outputs=[],
        )
